import sys

from PySide6.QtCore import QSize, Qt, QRect, Signal
from PySide6.QtGui import QPainter, QPen, QBrush, QColor
from PySide6.QtWidgets import QApplication, QMainWindow, QWidget, QDial, QVBoxLayout, QSizePolicy

class _Bar(QWidget):
    clickedValue = Signal(int)
    def __init__(self, steps):
        super().__init__()

        self.setSizePolicy(QSizePolicy.MinimumExpanding, QSizePolicy.MinimumExpanding)

        if isinstance(steps, list):
            self.n_steps = len(steps)
            self.steps = steps
            print("self.n_steps: ", self.n_steps)


        elif isinstance(steps, int):
            print("int")
            self.n_steps = steps
            self.steps = ["red"] * steps

        else:
            raise TypeError("steps must be a list or int")

        self._bar_solid_percent = 0.8
        self._background_color = QColor('black')
        self._padding = 4

    def _calculate_clicked_value(self, e):
        parent = self.parent()._dial
        vmin, vmax = parent.minimum(), parent.maximum()
        # 获取 Bar 控件高度
        d_height = self.size().height() + (self._padding * 2)
        # 计算每个 Bar 条目的高度
        step_size = d_height / self.n_steps
        # 计算点击位置
        click_y = e.y() - self._padding - step_size / 2

        # 计算点击位置占总高度的百分比
        pc = (d_height - click_y) / d_height
        # 根据百分比计算出 dial 对应的值
        value = int(vmin + pc * (vmax - vmin))
        # 通知 dial 控件更新
        self.clickedValue.emit(value)

    def mouseMoveEvent(self, event):
        self._calculate_clicked_value(event)
    
    def mousePressEvent(self, event):
        self._calculate_clicked_value(event)


    def sizeHint(self):
        return QSize(40, 120)

    def paintEvent(self, e):
        painter = QPainter(self)
        brush = QBrush()
        brush.setColor(self._background_color)
        brush.setStyle(Qt.SolidPattern)
        rect = QRect(0, 0, painter.device().width(), painter.device().height())
        painter.fillRect(rect, brush)

        # Get current state.
        dial = self.parent()._dial
        vmin, vmax = dial.minimum(), dial.maximum()
        value = dial.value()

        pen = painter.pen()
        pen.setColor(QColor('red'))
        painter.setPen(pen)

        font = painter.font()
        font.setFamily('Times')
        font.setPointSize(18)
        painter.setFont(font)

        pc = (value - vmin) / (vmax - vmin)
        # 计算需要绘制的 Bar 条目数量
        n_steps_to_draw = int(pc * self.n_steps)
        # painter.drawText(25, 25, "{}".format(n_steps_to_draw))
        d_height = painter.device().height() - (2 * self._padding)
        d_width = painter.device().width() - (2 * self._padding)

        # 总共能绘制多少个 Bar 条目
        step_size = d_height / self.n_steps
        # Bar 条目高度，这里乘一个百分比系数用以调节 Bar 条目高度
        bar_height = step_size * self._bar_solid_percent
        for n in range(n_steps_to_draw):
            # 设置每个 Bar 条目的背景颜色
            brush.setColor(QColor(self.steps[n]))
            # 计算每个 Bar 条目的位置
            ypos = (1 + n) * step_size
            # 绘制每个 Bar 条目
            rect = QRect(self._padding, self._padding * 2 + d_height - int(ypos), d_width, int(bar_height))
            painter.fillRect(rect, brush)

        painter.end()

    def _trigger_refresh(self):
        self.update()

class PowerBar(QWidget):
    def __init__(self, parent=None, steps=5):
        super().__init__(parent)

        layout = QVBoxLayout()
        self._bar = _Bar(steps)
        layout.addWidget(self._bar)
        
        self._dial = QDial()
        self._dial.valueChanged.connect(self._bar._trigger_refresh)
        layout.addWidget(self._dial)

        self._bar.clickedValue.connect(self._dial.setValue)
        self.setLayout(layout)

    def __getattr__(self, name):
        if name in self.__dict__:
            return self[name]
        
        try:
            return getattr(self._dial, name)
        except AttributeError: 
            raise AttributeError("'{}' object has no attribute '{}'".format(self._class__.__name__, name))
        
    def setColor(self, color):
        self._bar.steps = [color] * self._bar.n_steps
        self._bar.update()

    def setColors(self, colors):
        self._bar.n_steps = len(colors)
        self._bar.steps = colors
        self._bar.update()

    def setBarPadding(self, i):
        self._bar._padding = i
        self._bar.update()

    def setBarSolidPercent(self, f):
        self._bar._bar_solid_percent = float(f)
        self._bar.update()

    def setBackgroundColor(self, color):
        self._bar._background_color = QColor(color)
        self._bar.update()

    def setNotchesVisible(self, b):
        return self._dial.setNotchesVisible(b)

app = QApplication(sys.argv)
bar = PowerBar(steps=["#49006a", "#7a0177", "#ae017e", "#dd3497", "#f768a1","#fa9fb5", "#fcc5c0", "#fde0dd", "#fff7f3"])
bar.setBarPadding(2)
bar.setBarSolidPercent(0.9)
bar.setBackgroundColor('gray')
bar.show()
app.exec()

